package algorithm;

import java.util.HashMap;
import java.util.Map;

class TrieNode {
    private final Map<Character, TrieNode> children = new HashMap<>();
    private boolean isEndOfWord;

    public Map<Character, TrieNode> getChildren() {
        return children;
    }

    public boolean isEndOfWord() {
        return isEndOfWord;
    }

    public void setEndOfWord(boolean endOfWord) {
        isEndOfWord = endOfWord;
    }
}

class Trie {
    private final TrieNode root;

    public Trie() {
        root = new TrieNode();
    }

    // 插入一个字符串到 Trie 中
    public void insert(String word) {
        TrieNode current = root;
        for (char c : word.toCharArray()) {
            current = current.getChildren().computeIfAbsent(c, k -> new TrieNode());
        }
        current.setEndOfWord(true);
    }

    // 搜索前缀或完整的单词
    public boolean search(String word) {
        TrieNode current = root;
        for (char c : word.toCharArray()) {
            if (!current.getChildren().containsKey(c)) {
                return false;
            }
            current = current.getChildren().get(c);
        }
        return current.isEndOfWord();
    }

    // 查找是否有以给定前缀开头的单词
    public boolean startsWith(String prefix) {
        TrieNode current = root;
        for (char c : prefix.toCharArray()) {
            if (!current.getChildren().containsKey(c)) {
                return false;
            }
            current = current.getChildren().get(c);
        }
        return true;
    }

    public static void main(String[] args) {
        Trie trie = new Trie();

        // 插入一些示例字符串
        trie.insert("apple");
        trie.insert("app");
        trie.insert("banana");
        trie.insert("bat");

        // 搜索完整的单词
        System.out.println(trie.search("apple")); // 输出: true
        System.out.println(trie.search("app"));   // 输出: true
        System.out.println(trie.search("bat"));   // 输出: true
        System.out.println(trie.search("batman"));// 输出: false

        // 模糊查找前缀
        System.out.println(trie.startsWith("ba")); // 输出: true
        System.out.println(trie.startsWith("ap")); // 输出: true
        System.out.println(trie.startsWith("be")); // 输出: false
    }
}
